#!/usr/bin/env python
# Licensed under the MIT license:
# http://www.opensource.org/licenses/mit-license.php 

"""
Class to provide baseline TF-IDF processing of tweet corpus (after preprocess).
"""

from preprocess import wordlist
from tfidf import tfidf

TFIDF_OUTPUT = 'idf_out'
STOPWORD_OUTPUT = 'stopword_out'
SCORE_OUTPUT = 'scores_out'


def GetTfIdf():
  """Load wordlist, add users to tfidf (each user=document), save output."""
  words = wordlist.Wordlist(skipread=True)
  words.Load()
  # process TfIdf, save corpus and stopwords
  tf = tfidf.TfIdf()
  for user in words.worddict.keys():
    tf.add_input_document(' '.join(words.worddict[user]))
  tf.save_corpus_to_file(TFIDF_OUTPUT, STOPWORD_OUTPUT,
                         STOPWORD_PERCENTAGE_THRESHOLD=0.005)  # may tweak this
  # reload for stopwords, calculate and save scores
  tf = tfidf.TfIdf(corpus_filename=TFIDF_OUTPUT,
                   stopword_filename=STOPWORD_OUTPUT)
  sorted_terms = sorted(tf.term_num_docs.items(), key=tfidf.itemgetter(1),
                        reverse=True)
  output_file = open(SCORE_OUTPUT, "w")
  for term, num_docs in sorted_terms:
    score = tf.get_idf(term)
    if score:
      output_file.write(term + ": " + str(score) + "\n")
  output_file.close()
